import os
import sys
import logging

from PySide6.QtCore import QSettings
from PySide6.QtWidgets import QApplication, QMainWindow, QPushButton, QVBoxLayout, QWidget, QCheckBox, QLabel, \
    QTextEdit, QHBoxLayout, QTextBrowser, QTabWidget, QFileDialog, QLineEdit, QComboBox, QRadioButton
from scapy.all import sniff, sendp, Raw, wrpcap
from threading import Thread, Event

# 你的后端函数
from hd import start_capture, stop_capture, send_packet, save_packets, set_log_path

class NetworkTool(QMainWindow):
    def __init__(self):
        super().__init__()
        self.settings = QSettings("NetworkTool", "Settings")
        self.packets = []  # 定义 packets 变量
        self.logger = logging.getLogger()  # 定义 logger 变量
        self.initUI()

    def initUI(self):
        self.setWindowTitle('网络工具')
        self.setGeometry(100, 100, 400, 400)

        # 创建中心widget
        central_widget = QWidget()
        self.setCentralWidget(central_widget)

        # 创建布局
        layout = QVBoxLayout()

        # 创建日志路径输入框
        log_path_layout = QHBoxLayout()
        self.log_path_label = QLabel('日志路径:')
        self.log_path_edit = QLineEdit()
        self.log_path_edit.setText(self.settings.value("log_path", "log"))  # 读取保存的路径或使用默认值
        self.log_path_button = QPushButton('浏览')
        self.log_path_button.clicked.connect(self.browse_log_path)
        log_path_layout.addWidget(self.log_path_label)
        log_path_layout.addWidget(self.log_path_edit)
        log_path_layout.addWidget(self.log_path_button)
        layout.addLayout(log_path_layout)

        # 创建保存数据包路径输入框
        save_path_layout = QHBoxLayout()
        self.save_path_label = QLabel('保存数据包路径:')
        self.save_path_edit = QLineEdit()
        self.save_path_edit.setText(self.settings.value("save_path", "data"))  # 读取保存的路径或使用默认值
        self.save_path_button = QPushButton('浏览')
        self.save_path_button.clicked.connect(self.browse_save_path)
        save_path_layout.addWidget(self.save_path_label)
        save_path_layout.addWidget(self.save_path_edit)
        save_path_layout.addWidget(self.save_path_button)
        layout.addLayout(save_path_layout)

        # 创建保存数据包单选框
        self.save_packets_radio = QRadioButton('保存抓取到的包')
        self.save_packets_radio.setChecked(True)  # 默认选中
        layout.addWidget(self.save_packets_radio)

        # 创建tab页
        tab_widget = QTabWidget()
        receive_tab = QWidget()
        send_tab = QWidget()
        tab_widget.addTab(receive_tab, "接收")
        tab_widget.addTab(send_tab, "发送")
        layout.addWidget(tab_widget)

        # 创建接收部分的布局
        receive_layout = QVBoxLayout()
        receive_source_layout = QHBoxLayout()
        self.receive_source_checkbox1 = QCheckBox("数传服务器")
        self.receive_source_checkbox2 = QCheckBox("本地")
        self.receive_source_checkbox3 = QCheckBox("网络")
        receive_source_layout.addWidget(QLabel('接收来源:'))
        receive_source_layout.addWidget(self.receive_source_checkbox1)
        receive_source_layout.addWidget(self.receive_source_checkbox2)
        receive_source_layout.addWidget(self.receive_source_checkbox3)
        receive_layout.addLayout(receive_source_layout)
        self.receive_button = QPushButton('开始接收')
        self.receive_button.clicked.connect(self.start_receive)
        self.stop_receive_button = QPushButton('停止接收')
        self.stop_receive_button.clicked.connect(self.stop_receive)
        self.stop_receive_button.setEnabled(False)
        self.clear_receive_button = QPushButton('清空接收提示')
        self.clear_receive_button.clicked.connect(self.clear_receive_display)
        receive_layout.addWidget(self.receive_button)
        receive_layout.addWidget(self.stop_receive_button)
        receive_layout.addWidget(self.clear_receive_button)
        self.packet_display = QTextBrowser()
        receive_layout.addWidget(self.packet_display)
        receive_tab.setLayout(receive_layout)

        # 创建发送部分的布局
        send_layout = QVBoxLayout()
        send_destination_layout = QHBoxLayout()
        self.send_destination_checkbox1 = QCheckBox("数传服务器")
        self.send_destination_checkbox2 = QCheckBox("本地")
        self.send_destination_checkbox3 = QCheckBox("网络")
        send_destination_layout.addWidget(QLabel('发送目的:'))
        send_destination_layout.addWidget(self.send_destination_checkbox1)
        send_destination_layout.addWidget(self.send_destination_checkbox2)
        send_destination_layout.addWidget(self.send_destination_checkbox3)
        send_layout.addLayout(send_destination_layout)
        self.send_protocol_combo = QComboBox()
        self.send_protocol_combo.addItems(["TCP", "UDP"])
        send_layout.addWidget(QLabel('协议:'))
        send_layout.addWidget(self.send_protocol_combo)
        self.send_data_edit = QTextEdit()
        self.send_data_edit.setPlaceholderText("输入要发送的数据")
        send_layout.addWidget(self.send_data_edit)
        self.add_file_button = QPushButton('添加文件')
        self.add_file_button.clicked.connect(self.add_file)
        send_layout.addWidget(self.add_file_button)
        self.send_button = QPushButton('发送')
        self.send_button.clicked.connect(self.send_packet)
        self.clear_send_button = QPushButton('清空发送提示')
        self.clear_send_button.clicked.connect(self.clear_send_display)
        send_layout.addWidget(self.send_button)
        send_layout.addWidget(self.clear_send_button)
        self.send_status_display = QTextBrowser()
        send_layout.addWidget(self.send_status_display)
        send_tab.setLayout(send_layout)

        central_widget.setLayout(layout)

    def browse_log_path(self):
        log_path = QFileDialog.getExistingDirectory(self, "选择日志保存路径")
        if log_path:
            self.log_path_edit.setText(log_path)
            self.settings.setValue("log_path", log_path)  # 保存路径
            set_log_path(log_path)  # 更新后端的日志路径

    def browse_save_path(self):
        save_path = QFileDialog.getExistingDirectory(self, "选择数据包保存路径")
        if save_path:
            self.save_path_edit.setText(save_path)
            self.settings.setValue("save_path", save_path)  # 保存路径

    def start_receive(self):
        sources = []
        if self.receive_source_checkbox1.isChecked():
            sources.append("1")
        if self.receive_source_checkbox2.isChecked():
            sources.append("2")
        if self.receive_source_checkbox3.isChecked():
            sources.append("3")
        filter_rule = "ip src 192.168.1.100 and ip dst 192.168.1.200 and (tcp port 80 or udp port 80)"
        log_path = self.log_path_edit.text()
        save_path = self.save_path_edit.text()
        start_capture(sources, filter_rule, log_path, self.packets)  # 传递 packets 变量
        self.receive_button.setEnabled(False)
        self.stop_receive_button.setEnabled(True)
        self.save_packets_radio.setEnabled(False)  # 置灰单选框
        self.receive_source_checkbox1.setEnabled(False)
        self.receive_source_checkbox2.setEnabled(False)
        self.receive_source_checkbox3.setEnabled(False)
        self.packet_display.append("开始接收数据包...")

    def stop_receive(self):
        stop_capture()
        self.receive_button.setEnabled(True)
        self.stop_receive_button.setEnabled(False)
        self.save_packets_radio.setEnabled(True)  # 恢复单选框可操作
        self.receive_source_checkbox1.setEnabled(True)
        self.receive_source_checkbox2.setEnabled(True)
        self.receive_source_checkbox3.setEnabled(True)
        if self.save_packets_radio.isChecked():
            save_path = self.save_path_edit.text()  # 获取保存路径
            self.save_packets(save_path)  # 调用类方法
            self.packet_display.append("抓取到的包已保存到文件中。")
        self.packet_display.append("停止接收数据包。")

    def clear_receive_display(self):
        self.packet_display.clear()

    def clear_send_display(self):
        self.send_status_display.clear()

    def send_packet(self):
        destinations = []
        if self.send_destination_checkbox1.isChecked():
            destinations.append("1")
        if self.send_destination_checkbox2.isChecked():
            destinations.append("2")
        if self.send_destination_checkbox3.isChecked():
            destinations.append("3")
        protocol = self.send_protocol_combo.currentText()
        data = self.send_data_edit.toPlainText().encode('utf-8')  # 将文本数据编码为二进制
        file_data = b""

        if hasattr(self, 'file_path'):
            try:
                with open(self.file_path, 'rb') as f:  # 以二进制模式读取文件
                    file_data = f.read()
                data += file_data  # 直接将文件的二进制内容追加到数据中
                self.send_status_display.append(f"文件 {self.file_path} 已添加。")
            except Exception as e:
                self.send_status_display.append(f"读取文件 {self.file_path} 时发生错误: {e}")
                return  # 退出函数，不发送数据

        log_path = self.log_path_edit.text()
        for destination in destinations:
            try:
                send_packet(destination, protocol, data, log_path)
                if hasattr(self, 'file_path'):
                    self.send_status_display.append(f"文件 {self.file_path} 发送成功: 本地 -> {destination}")
                else:
                    self.send_status_display.append(f"数据发送成功: 本地 -> {destination}")
            except Exception as e:
                self.send_status_display.append(f"发送失败: 本地 -> {destination}, 错误信息: {e}")
    def add_file(self):
        self.file_path, _ = QFileDialog.getOpenFileName(self, "选择文件", "", "All Files (*)")
        if self.file_path:
            self.send_status_display.append(f"文件 {self.file_path} 已添加。")

    def save_packets(self, save_path):
        # 调用后端的 save_packets 函数来保存抓取到的数据包
        save_packets(save_path, self.packets)
if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = NetworkTool()
    ex.show()
    sys.exit(app.exec())